#!/usr/bin/env python
import sys
import math
import os, time, random
from multiprocessing import Pool

def process_images(name, task_cmd):
    print("=============== Checking env =============")
    print("[SUBJECTS_DIR]: %s" % (os.getenv('SUBJECTS_DIR')))
    print("================== Done ==================")

    print('Run task %s (%s)...' % (name, os.getpid()))
    start = time.time()
    print('[RUN CMD]: %s' % (task_cmd))
    os.system(task_cmd)
    end = time.time()
    print('Task %s runs %0.2f seconds.' % (name, (end - start)))

def get_task_infos(src_task_file, processed_rest_dir, processed_t1_dir, reg_stem = '_rest_skip10_stc_mc'):
    cmd_list = {}

    fp = open(src_task_file, "r")
    for line in fp.readlines():
        subj_name = line.replace("\n", "")

        sess_list = []
        for run_folder in ['001', '002', '003']:
            sud_dir = os.path.join(subj_name, 'bold')
            sud_dir = os.path.join(sud_dir, run_folder)
            sud_dir = os.path.join(processed_rest_dir, sud_dir)

            tkreg_cmd = "tkregister2"
            # --s
            tkreg_cmd = tkreg_cmd + " --s " + subj_name
            # --sd
            tkreg_cmd = tkreg_cmd + " --sd " + processed_t1_dir
            # --mov
            nii_file = subj_name + '_bld' + run_folder + reg_stem + '.nii.gz'
            full_nii_file = os.path.join(sud_dir, nii_file)
            tkreg_cmd = tkreg_cmd + " --mov " + full_nii_file
            # --reg
            reg_file = 'init-fsl/' + subj_name + '_bld' + run_folder + reg_stem + '_reg.dat'
            full_reg_file = os.path.join(sud_dir, reg_file)
            tkreg_cmd = tkreg_cmd + " --reg " + full_reg_file
            # --surf
            tkreg_cmd = tkreg_cmd + " --surf"

            sess_list.append(tkreg_cmd)

        cmd_list[subj_name] = sess_list

    return cmd_list, len(cmd_list)

def set_process_env():
    # set env
    os.environ["SUBJECTS_DIR"] = "/mri_projects/ASD_Analysis/DataProcessed/T1"

def screenshot(trg_dir, sub_name, sess_id):
    print("# Shot Screen... #")
    output_file = sub_name + '_sess_' + str(sess_id) + '.png'

    output_path = os.path.join(trg_dir, output_file)
    shot_cmd = 'shutter -s=650,450,400,400 -e -o ' + output_path
    print("[CMD]: %s" % (shot_cmd))

    os.system(shot_cmd)

if __name__=='__main__':
    print('Parent process %s.' % os.getpid())

    set_process_env()

    project_dir = '/mri_projects/ASD_Analysis'
    symbol      = "OHSU"
    sub_type    = "HC"
    list_file   = 'variability_list_' + sub_type + '_' + symbol + '.txt'
    reg_stem    = '_rest_skip10_stc_mc'

    src_task_file = os.path.join(project_dir, os.path.join('Lists/', list_file))
    processed_rest_dir = os.path.join(project_dir, 'DataProcessed/rest/')
    processed_t1_dir = os.path.join(project_dir, 'DataProcessed/T1/')

    task_dict, task_size = get_task_infos(src_task_file, processed_rest_dir, processed_t1_dir, reg_stem)

    output_dir = os.path.join(project_dir, "Figures/QC/{}/{}".format(symbol, sub_type))
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    process_num = 1
    p = Pool(process_num)

    for (name, cmd_list) in task_dict.items():
        print("Processing %s" % (name))
        for i in range(len(cmd_list)):
            p.apply_async(process_images, args=(i, cmd_list[i],))
            # p.apply_async(screenshot, args=(output_dir, name, i + 1, ))

    print('Waiting for all subprocesses done...')
    p.close()
    p.join()
